package org.zjx.server;

import lombok.SneakyThrows;
import org.zjx.core.SocketConfig;
import org.zjx.handler.CommandHandlerFactory;
import org.zjx.message.IMessage;
import org.zjx.message.KryoSerializer;
import org.zjx.message.MessageFactory;
import org.zjx.message.messages.*;

import java.io.*;
import java.lang.reflect.Method;
import java.net.Socket;
import java.util.Map;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.zjx.server.scanner.ServiceScanner;

public class SocketServer extends BaseSocketServer {
    private static final Logger logger = LoggerFactory.getLogger(SocketServer.class);

    private final ScheduledExecutorService heartbeatChecker;
    private final ServiceRegistry serviceRegistry;

    public SocketServer(SocketConfig config) {
        super(config);
        this.heartbeatChecker = Executors.newScheduledThreadPool(1);

        // 自动扫描并注册服务
        this.serviceRegistry = new ServiceRegistry();
        Map<Class<?>, Object> services = ServiceScanner.scanServices();
        services.forEach((serviceInterface, implementation) -> {
            serviceRegistry.registerService(serviceInterface, implementation);
            logger.info("Auto-registered service: {}", serviceInterface.getName());
        });
    }

    @Override
    protected void handleConnection(Socket clientSocket) {
        new Thread(new ClientHandler(clientSocket)).start();
    }

    @Override
    public void close() throws IOException {
        // 关闭心跳检测线程池
        if (heartbeatChecker != null) {
            heartbeatChecker.shutdownNow();
            try {
                if (!heartbeatChecker.awaitTermination(3, TimeUnit.SECONDS)) {
                    logger.warn("Heartbeat checker did not terminate in time");
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }
        super.close();
    }

    private class ClientHandler implements Runnable {
        private final Logger handlerLogger = LoggerFactory.getLogger(ClientHandler.class);
        private final Socket socket;
        private final DataOutputStream dataOut;
        private final DataInputStream dataIn;
        private final long lastActivity = System.currentTimeMillis();
        private volatile boolean connectionActive = true;
        private ScheduledFuture<?> heartbeatFuture;

        @SneakyThrows
        public ClientHandler(Socket socket) {
            this.socket = socket;
            this.dataOut = new DataOutputStream(socket.getOutputStream());
            this.dataIn = new DataInputStream(socket.getInputStream());
        }

        @Override
        public void run() {
            // 启动心跳检测（如果配置）
            if (config.getHeartbeatTimeout() > 0) {
                heartbeatFuture = startHeartbeatMonitor();
            }

            try {
                processClientRequests();
            } catch (IOException e) {
                if (running && connectionActive) {
                    logger.error(e.getMessage());
                }
            } finally {
                cleanupConnection();
            }
        }

        private void processClientRequests() throws IOException {
            while (running && connectionActive) {
                int length = dataIn.readInt();       // 读长度头
                byte[] data = new byte[length];
                dataIn.readFully(data);              // 读完整数据
                IMessage message = MessageFactory.deserialize(data);
                processMessage(message);
            }
        }

        private void processMessage(IMessage message) {
            if (message instanceof CommandMessage) {
                handleCommand((CommandMessage) message);
            } else if (message instanceof RegisterWorkerMessage) {
                handleRegisterRoute((RegisterWorkerMessage) message);
            } else if (message instanceof RegisterClientMessage) {
                handleRegisterClient((RegisterClientMessage) message);
            } else if (message instanceof HeartbeatMessage) {
                handleHeartbeat((HeartbeatMessage) message);
            } else if (message instanceof RpcRequest) {
                handleRpc((RpcRequest) message);
            } else {
                sendError();
            }
        }

        // 更新心跳时间并回复
        private void handleHeartbeat(HeartbeatMessage message) {
            handlerLogger.debug("Received heartbeat from client");
        }

        //注册路由
        private void handleRegisterRoute(RegisterWorkerMessage routeMessage) {
            handlerLogger.info("Worker registration received: {}", routeMessage.getWorkerId());
        }

        //注册客户端
        private void handleRegisterClient(RegisterClientMessage message) {
            handlerLogger.info("Client registration received");
        }

        //RPC代理
        @SneakyThrows
        private void handleRpc(RpcRequest request) {
            handlerLogger.debug("Handling RPC: {}.{}", request.getClassName(), request.getMethodName());
            // 处理请求
            Object service = serviceRegistry.getService(request.getClassName());
            if (service == null) {
                throw new ClassNotFoundException(request.getClassName() + " not found");
            }

            Method method = service.getClass().getMethod(
                    request.getMethodName(), request.getParameterTypes());
            Object result = method.invoke(service, request.getParameters());

            // 返回响应
            RpcResponse response = new RpcResponse();
            response.setMessageId(request.getMessageId());
            response.setResult(result);
            sendMessage(response);
        }

        // 普通命令处理
        private void handleCommand(CommandMessage cmd) {
            ResponseMessage response = executeLocally(cmd);
            sendMessage(response);
        }

        private ResponseMessage executeLocally(CommandMessage command) {
            return CommandHandlerFactory.getHandler(command.getCommand())
                    .map(handler -> handler.handle(command))
                    .orElse(ResponseMessage.error(command, "Unsupported command"));
        }

        private ScheduledFuture<?> startHeartbeatMonitor() {
            return heartbeatChecker.scheduleAtFixedRate(() -> {
                // 双重检查确保连接有效
                if (!connectionActive || socket.isClosed()) {
                    return;
                }

                long inactiveTime = System.currentTimeMillis() - lastActivity;
                if (inactiveTime > config.getHeartbeatTimeout()) {
                    handlerLogger.warn("Client heartbeat timeout");

                    // 原子性关闭操作
                    synchronized (this) {
                        if (connectionActive) {
                            connectionActive = false;
                            closeQuietly(socket);
                        }
                    }
                }
            }, 10, 10, TimeUnit.SECONDS);
        }

        private void cleanupConnection() {
            // 取消心跳检测任务
            if (heartbeatFuture != null) {
                heartbeatFuture.cancel(true);
            }

            // 标记连接不活跃
            connectionActive = false;

            // 安全关闭资源
            closeQuietly(dataIn);
            closeQuietly(dataOut);
            closeQuietly(socket);

            logger.info("Client disconnected");
        }

        private void sendError() {
            sendMessage(ResponseMessage.error(null, "接收到未知消息类型"));
        }

        @SneakyThrows
        private void sendMessage(IMessage message) {
            if (connectionActive) {
                byte[] data = MessageFactory.serialize(message);
                dataOut.writeInt(data.length);  // 写长度头
                dataOut.write(data);            // 写序列化数据
                dataOut.flush();
            }
        }
    }

}
